from common.config import D, T, G, get_config_path
from engine.strategy import Strategy
from engine.backtest import Backtest
import importlib
import os
import pandas as pd

class TaskRunner:
    def run(self, config: dict):
        df = self.prepare_data(config)
        df.to_csv(get_config_path() + '/cache.csv')
        s = self.prepare_strategy(config)
        #print(s)
        b = Backtest(df=df)
        df = b.run(s)
        df.to_csv(get_config_path() + '/results/{}.csv'.format(s.name))


    def prepare_data(self,config:dict):
        #df = D.load(config['codes'], start_time=config['date_start'], end_time=config['date_end'])

        fields = []
        if 'rules' in config.keys():
            rules = config['rules']
            buy_rules = rules['buy_rules']
            sell_rules = rules['sell_rules']

            fields = buy_rules + sell_rules

        if 'features' in config.keys():
            fields += config['features']

        df = D.load(config['codes'], start_time=config['date_start'], end_time=config['date_end'], fields=fields, names=fields)

        if 'rules' in config.keys():
            rules = config['rules']
            buy_temp = df[rules['buy_rules']].sum(axis=1)
            df['to_buy'] = buy_temp >= rules['buy_min_count']
            sell_temp = df[rules['sell_rules']].sum(axis=1)
            df['to_sell'] = sell_temp >= rules['sell_min_count']
        return df

    def prepare_strategy(self, config:dict):
        algos = config['algos']
        algo_list = [

        ]
        for algo_dict in algos:
            module = importlib.import_module(algo_dict['module'])
            kwargs = None

            algo_class = getattr(module,algo_dict['class'])

            if 'kwargs' in algo_dict.keys():
                kwargs = algo_dict['kwargs']
                algo_obj = algo_class(**kwargs)
            else:
                algo_obj = algo_class()

            algo_list.append(algo_obj)

        s = Strategy(name=config['name'], algo_list=algo_list)
        return s

    @staticmethod
    def analysis(start, end, benchmarks=[], task_names=[]):
        equities = []
        for benchmark in benchmarks:
            bench_df = D.get_code_df(benchmark)[start:end]
            se = (bench_df['rate']+1).cumprod()
            se.name = benchmark
            equities.append(se)

        for name in task_names:
            filename = get_config_path()+'/results/{}.csv'.format(name)
            if os.path.exists(filename):
                df = pd.read_csv(filename)
                df['date'] = df['date'].apply(lambda x:str(x))
                df.index = df['date']
                se =(df['rate']+1).cumprod()
                se.name = name
                equities.append(se)

        df_equities = pd.concat(equities,axis=1)
        df_equities.dropna(inplace=True)
        print(df_equities)

        from engine.performance import PerformanceUtils

        df_ratios, df_corr, df_years = PerformanceUtils().calc_equity(df_equity=df_equities)
        G.notify(G.GRAPH_RESULT,
                 {
                     'equity': df_equities,
                     'ratio': df_ratios,
                     'corr': df_corr,
                     'yearly': df_years

                 })
        return df_equities, df_ratios, df_corr, df_years


if __name__ == '__main__':
    task_config = T.get_task('RSRS指标轮动')
    print(task_config)
    TaskRunner().run(task_config)